
import sys
import time

from gxl_ai_utils.utils import utils_file
# from infer_runtime import get_feat_from_wav_path
sys.path.insert(0, '../../../../')
sys.path.insert(0, '.')

try:
    from wenet.utils.init_tokenizer import init_tokenizer
    from gxl_ai_utils.config.gxl_config import GxlNode
    from wenet.utils.init_model import init_model
    import logging
    import librosa
    import torch
    import torchaudio
except ImportError:
    pass
is_npu = True

try:
    import torch_npu
except ImportError:
    is_npu = False
    print("torch_npu is not available. if you want to use npu, please install it.")

def do_resample(input_wav_path, output_wav_path):
    """"""
    waveform, sample_rate = torchaudio.load(input_wav_path)
    # 检查音频的维度
    num_channels = waveform.shape[0]
    # 如果音频是多通道的，则进行通道平均
    if num_channels > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    waveform = torchaudio.transforms.Resample(
        orig_freq=sample_rate, new_freq=16000)(waveform)
    utils_file.makedir_for_file(output_wav_path)
    torchaudio.save(output_wav_path, waveform, 16000)
def get_feat_from_wav_path(input_wav_path, device):
    """
    获取音频的特征
    Args:
        input_wav_path: str

    Returns:
        feat: tensor, shape=(1, T, 80)
        feat_lens: tensor, shape=(1,)
    """
    timestamp_ms = int(time.time() * 1000)
    now_file_tmp_path_resample = f'~/.cache/.temp/{timestamp_ms}_resample.wav'
    do_resample(input_wav_path, now_file_tmp_path_resample)
    input_wav_path = now_file_tmp_path_resample
    waveform, sample_rate = torchaudio.load(input_wav_path)
    waveform = waveform.squeeze(0)  # (channel=1, sample) -> (sample,)
    window = torch.hann_window(400)
    stft = torch.stft(waveform,
                      400,
                      160,
                      window=window,
                      return_complex=True)
    magnitudes = stft[..., :-1].abs() ** 2

    filters = torch.from_numpy(
        librosa.filters.mel(sr=sample_rate,
                            n_fft=400,
                            n_mels=80))
    mel_spec = filters @ magnitudes

    # NOTE(xcsong): https://github.com/openai/whisper/discussions/269
    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0
    feat = log_spec.transpose(0, 1)
    feat_lens = torch.tensor([feat.shape[0]], dtype=torch.int64).to(device)
    feat = feat.unsqueeze(0).to(device)
    return feat, feat_lens


def init_model_my(gpu_id=0, checkpoint_path=None, config_path=None):
    logging.basicConfig(level=logging.DEBUG,
                        format='%(asctime)s %(levelname)s %(message)s')
    if checkpoint_path is None:
        checkpoint_path = "/mnt/sfs/asr/code/osum_xlgeng_3B/examples/wenetspeech/whisper/exp/epoch_3_7_LLM3Binstruct_5Ws2t/step_9999.pt"
    if config_path is None:
        config_path = "../conf/config_llm_huawei_instruct_3B_cosyvoice1-token.yaml"
    args = GxlNode({
        "checkpoint": checkpoint_path,
    })
    configs = utils_file.load_dict_from_yaml(config_path)
    model, configs = init_model(args, configs)
    if is_npu:
        device = torch.device(f'npu:{gpu_id}')
    else:
        device =torch.device(f'cuda:{gpu_id}')
    model = model.to(device)
    model.eval()
    tokenizer = init_tokenizer(configs)
    print(model)
    return model, tokenizer, device


# model, tokenizer, device = init_model_my()

def do_TTS(model, device, text_for_tts):
    input_prompt = "恳请将如下文本转换为其对应的语音token，力求生成最为流畅、自然的语音。"
    res_text = model.generate_tts(device=device, text=text_for_tts, prompt=input_prompt)[0]
    return res_text

def do_chat_text2text(model, device, text_for_chat):
    if isinstance(device, str):
        device = torch.device(f"npu:{int(device)}")
    res_text = model.generate_text2text(device=device, text=text_for_chat)[0]
    return res_text

def do_chat_s2s(model, device, wav_path):
    """"""
    prompt = "先根据语音输入，直接以文字形式进行回答或对话，接着再生成语音token。"
    if isinstance(device, str):
        device = torch.device(f"npu:{int(device)}")
    feat, feat_lens = get_feat_from_wav_path(wav_path, device)
    feat = feat.to(device)
    feat_lens = feat_lens.to(device)
    res_text = model.generate_s2s(wavs=feat, wavs_len=feat_lens, prompt=prompt)[0]
    return res_text